package com.github.common.extend.demo.config.filetask.excel;

import com.alibaba.excel.support.ExcelTypeEnum;
import com.alibaba.fastjson.JSON;
import com.github.common.extend.demo.config.filetask.*;
import com.github.common.extend.demo.config.filetask.download.FileTaskDownloader;
import com.github.common.extend.demo.util.excel.XExcelUtils;
import com.github.common.extend.demo.util.iterator.partition.fixed.LocalDateTimePartitionParam;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.http.HttpHeaders;
import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import java.io.File;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.time.LocalDateTime;
import java.time.temporal.TemporalUnit;
import java.util.Iterator;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import java.util.function.Supplier;

import static org.springframework.http.HttpHeaders.CONTENT_DISPOSITION;
import static org.springframework.http.MediaType.APPLICATION_OCTET_STREAM_VALUE;

/**
 * @author wangjj7
 * @date 2023/8/29
 * @description
 */
@Slf4j
public class ExcelTaskDownloaderImpl implements ExcelTaskDownloader {

    private final FileTaskDownloader fileTaskDownloader;

    private final Supplier<String> idGenerator;

    private final ExecutorService executor;

    private static final int TIMEOUT_SECOND = 55;

    public ExcelTaskDownloaderImpl(FileTaskDownloader fileTaskDownloader,
                                   Supplier<String> idGenerator,
                                   ExecutorService executor)
    {
        Objects.requireNonNull(fileTaskDownloader);
        Objects.requireNonNull(idGenerator);
        Objects.requireNonNull(executor);

        this.fileTaskDownloader = fileTaskDownloader;
        this.idGenerator = idGenerator;
        this.executor = executor;
    }

    public static File defaultTmpWriteDir(){
        return new File(System.getProperty("java.io.tmpdir") + File.separator + "fsb");
    }

    public static File defaultTmpWriteFile(String fileName){
        return new File(defaultTmpWriteDir().getPath() + File.separator + fileName);
    }

    @Override
    public <T> void writeThenDownload(Class<T> headType, Iterator<T> iterator,
                                      String fileName, HttpServletResponse response){
        try(ServletOutputStream outputStream = response.getOutputStream()) {
            if (!fileName.contains(".")){
                fileName = fileName + ExcelTypeEnum.XLSX.getValue();
            }
            String encodeFileName = fileName;
            try {
                encodeFileName = URLEncoder.encode(fileName,"utf-8");
            } catch (UnsupportedEncodingException e) {
                log.warn("文件名转码失败" ,e.getMessage());
            }

            response.setContentType(APPLICATION_OCTET_STREAM_VALUE);
            response.setHeader(CONTENT_DISPOSITION, "attachment;filename=" + encodeFileName);
            response.setHeader(HttpHeaders.ACCESS_CONTROL_EXPOSE_HEADERS, CONTENT_DISPOSITION);
            XExcelUtils.write(headType ,outputStream ,iterator);
        } catch (IOException e) {
            throw new IllegalStateException(e);
        }
    }

    @Override
    public <T> void submitWriteTaskThenDownload(FileTaskParam fileTaskParam, Class<T> headType,
                                                Iterator<T> iterator, HttpServletResponse response)
    {
        String id = fileTaskParam.getId();
        if (StringUtils.isBlank(id)){
            id = String.valueOf(idGenerator.get());
            fileTaskParam.setId(id);
        }
        File tmpWriteFile = defaultTmpWriteFile(id + ExcelTypeEnum.XLSX.getValue());

        FileSourceBuilder excelBuilder = new FileSourceBuilder(tmpWriteFile ,
                (fileOutputStream ,rateRefresh) -> XExcelUtils.write(headType, fileOutputStream, iterator)
        );
        fileTaskParam.setDataSourceBuilder(excelBuilder);
        try {
            fileTaskDownloader.submitTaskThenDownload(fileTaskParam ,response);
        }finally {
            excelBuilder.clear();
        }
    }

    @Override
    public <T> void asyncSubmitWriteTaskThenDownload(FileTaskParam fileTaskParam, Class<T> headType,
                                                     Iterator<T> iterator, HttpServletResponse response,
                                                     FileTaskCallback fileTaskCallback)
    {
        String id = fileTaskParam.getId();
        if (StringUtils.isBlank(id)){
            id = String.valueOf(idGenerator.get());
            fileTaskParam.setId(id);
        }
        File tmpWriteFile = defaultTmpWriteFile(id + ExcelTypeEnum.XLSX.getValue());

        FileSourceBuilder excelBuilder = new FileSourceBuilder(tmpWriteFile ,(fileOutputStream ,rateRefresh) -> {
            XExcelUtils.write(headType ,fileOutputStream, iterator);
        });
        fileTaskParam.setDataSourceBuilder(excelBuilder);

        fileTaskDownloader.asyncSubmitTaskThenDownload(
                fileTaskParam, fileClearCallback(fileTaskCallback ,excelBuilder),
                response, TIMEOUT_SECOND, TimeUnit.SECONDS);//前端一般是60秒，这里设置55秒吧
    }

    @Override
    public <T ,P extends LocalDateTimePartitionParam<P>> void submitWriteByTimePartitionTaskThenDownload(
            FileTaskParam fileTaskParam, P query, Class<T> headType,
            int partitionSize, int partitionInterval, TemporalUnit unit, boolean asc,
            Function<P, List<T>> queryFunction, HttpServletResponse response)
    {
        FileSourceBuilder excelBuilder = initDataSourceBuilder(fileTaskParam, query, headType, partitionSize,
                partitionInterval, unit, asc, queryFunction);
        try {
            fileTaskDownloader.submitTaskThenDownload(fileTaskParam ,response);
        }finally {
            excelBuilder.clear();
        }
    }

    @Override
    public <T ,P extends LocalDateTimePartitionParam<P>> void asyncSubmitWriteByTimePartitionTaskThenDownload(
            FileTaskParam fileTaskParam, P query, Class<T> headType,
            int partitionSize, int partitionInterval, TemporalUnit unit, boolean asc,
            Function<P, List<T>> queryFunction, HttpServletResponse response,
            FileTaskCallback fileTaskCallback)
    {
        FileSourceBuilder excelBuilder = initDataSourceBuilder(fileTaskParam, query, headType, partitionSize,
                partitionInterval, unit, asc, queryFunction);

        fileTaskDownloader.asyncSubmitTaskThenDownload(
                fileTaskParam, fileClearCallback(fileTaskCallback ,excelBuilder),
                response, 55, TimeUnit.SECONDS);//前端一般是60秒，这里设置55秒吧
    }

    private <T ,P extends LocalDateTimePartitionParam<P>> FileSourceBuilder initDataSourceBuilder(
            FileTaskParam fileTaskParam, P query, Class<T> headType,
            int partitionSize, int partitionInterval, TemporalUnit intervalUnit, boolean asc,
            Function<P, List<T>> queryFunction)
    {
        String id = fileTaskParam.getId();
        if (StringUtils.isBlank(id)){
            id = String.valueOf(idGenerator.get());
            fileTaskParam.setId(id);
        }
        File tmpWriteFile = defaultTmpWriteFile(id + ExcelTypeEnum.XLSX.getValue());

        LocalDateTime startLocalTime = query.$readStartTime();
        LocalDateTime endLocalTime = query.$readEndTime();
        long total = intervalUnit.between(startLocalTime ,endLocalTime);
        int partitionCount = (int) total / partitionSize;
        AtomicInteger rate = new AtomicInteger();

        //大文件直接写到内存中可能会导致OOM，先写到文件里
        String finalId = id;
        FileSourceBuilder excelBuilder = new FileSourceBuilder(tmpWriteFile ,
                (fileOutputStream ,rateRefresh) -> XExcelUtils.writeByTimePartition(fileOutputStream ,query ,
                        headType, partitionSize, partitionInterval, intervalUnit, asc,
                        partitionQuery -> {
                            log.info("任务【{}】时间分片【{}->{}】查询，分片参数【{}】" ,
                                    finalId,partitionQuery.$readStartTime() ,partitionQuery.$readEndTime() ,
                                    JSON.toJSON(partitionQuery));
                            if (rateRefresh != null) {
                                rateRefresh.accept(100d * rate.getAndIncrement() / partitionCount);
                            }
                            List<T> results = queryFunction.apply(partitionQuery);
                            log.info("任务【{}】时间分片【{}->{}】查询，数量【{}】" ,
                                    finalId,partitionQuery.$readStartTime() ,partitionQuery.$readEndTime() ,
                                    results.size());
                            return results;
                        } ,executor)
        );

        fileTaskParam.setDataSourceBuilder(excelBuilder);
        return excelBuilder;
    }

    private FileTaskCallback fileClearCallback(FileTaskCallback fileTaskCallback ,FileSourceBuilder fileSourceBuilder)
    {
        return new FileTaskCallback() {
            @Override
            public void preExe(FileTaskParam taskSubmitDTO) {
                if (fileTaskCallback != null){
                    fileTaskCallback.preExe(taskSubmitDTO);
                }
            }

            @Override
            public void postExe(FileTask fileTask, Exception e) {
                if (fileTaskCallback != null){
                    try {
                        fileTaskCallback.postExe(fileTask ,e);
                    }finally {
                        //确保删除临时文件
                        fileSourceBuilder.clear();
                    }
                }else {
                    fileSourceBuilder.clear();
                }
            }
        };
    }

}
